"""
MCP WebSocket client — 兼容服务端并修复常见超时原因

改动要点
- 发送请求时使用数字 id（未加引号），这是大多数 JSON-RPC 示例（包括你服务端示例）使用的形式。
- 为兼容服务器可能以字符串或数字返回 id，pending 映射同时保存两个键：int id 和 str(id)。
- 接收时按多种形式匹配 pending（原始、int、str）。
- 详细记录原始发送/接收帧，便于排查服务器是否真的返回了响应。
- 提供 async context manager 接口 (async with)。

用法
    python3 mcp_ws_client_fix.py ws://localhost:8000/mcp

或者在代码中：
    async with MCPWebSocketClient(uri) as client:
        result = await client.request_and_get_result("hello", {"name":"world"})
"""
from __future__ import annotations
import asyncio
import json
import logging
import itertools
from typing import Any, Dict, Optional, Callable, Awaitable

import websockets
from websockets.legacy.client import WebSocketClientProtocol

logger = logging.getLogger("mcp_ws_client_fix")
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
logger.addHandler(handler)
logger.setLevel(logging.DEBUG)

class MCPWebSocketClient:
    def __init__(
        self,
        uri: str,
        on_notification: Optional[Callable[[Any], Awaitable[None]]] = None,
        timeout: float = 10.0,
    ):
        self.uri = uri
        self.on_notification = on_notification
        self.timeout = timeout

        self._ws: Optional[WebSocketClientProtocol] = None
        self._reader_task: Optional[asyncio.Task] = None

        # pending: map keys can be int or str -> future
        self._pending: Dict[Any, asyncio.Future] = {}
        self._id_counter = itertools.count(1)
        self._lock = asyncio.Lock()
        self._closed = True

    async def __aenter__(self):
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self.close()

    def _is_connected(self) -> bool:
        if self._ws is None:
            return False
        closed = getattr(self._ws, "closed", None)
        if isinstance(closed, bool):
            return not closed
        open_attr = getattr(self._ws, "open", None)
        if isinstance(open_attr, bool):
            return open_attr
        close_code = getattr(self._ws, "close_code", None)
        if close_code is None:
            return True
        return False

    async def connect(self) -> None:
        async with self._lock:
            if self._ws is not None and self._is_connected():
                return
            logger.info("Connecting to %s", self.uri)
            self._ws = await websockets.connect(self.uri, subprotocols=["mcp"])
            logger.info("Connected (accepted subprotocol=%r)", getattr(self._ws, "subprotocol", None))
            self._reader_task = asyncio.create_task(self._reader_loop())
            self._closed = False

    async def close(self) -> None:
        async with self._lock:
            if self._closed:
                return
            self._closed = True
            if self._reader_task:
                self._reader_task.cancel()
                self._reader_task = None
            if self._ws:
                try:
                    close_coro = getattr(self._ws, "close", None)
                    if callable(close_coro):
                        await close_coro()
                except Exception:
                    logger.exception("error closing websocket")
                self._ws = None
            for fut in list(self._pending.values()):
                if not fut.done():
                    fut.set_exception(ConnectionError("connection closed"))
            self._pending.clear()
            logger.info("Closed")

    async def _reader_loop(self):
        if self._ws is None:
            return
        try:
            async for text in self._ws:
                logger.debug("RECV RAW: %s", text)
                try:
                    data = json.loads(text)
                except Exception:
                    logger.exception("failed to parse incoming JSON: %r", text)
                    continue

                msg_id = data.get("id")
                if msg_id is not None:
                    fut = None
                    # 1) try direct lookup
                    fut = self._pending.pop(msg_id, None)
                    # 2) if msg_id is string digits, try int
                    if fut is None and isinstance(msg_id, str) and msg_id.isdigit():
                        try:
                            fut = self._pending.pop(int(msg_id), None)
                        except Exception:
                            fut = None
                    # 3) try str(msg_id)
                    if fut is None:
                        try:
                            fut = self._pending.pop(str(msg_id), None)
                        except Exception:
                            fut = None

                    if fut:
                        if not fut.done():
                            fut.set_result(data)
                        else:
                            logger.debug("future already done for id %r", msg_id)
                    else:
                        logger.warning("received response for unknown id %r (payload: %r)", msg_id, data)
                    continue

                # Notification
                logger.info("notification received: %r", data)
                if self.on_notification:
                    try:
                        asyncio.create_task(self.on_notification(data))
                    except Exception:
                        logger.exception("on_notification handler raised")
        except asyncio.CancelledError:
            logger.debug("reader loop cancelled")
        except Exception:
            logger.exception("reader loop error")
        finally:
            await self.close()

    async def send_request(self, method: str, params: Optional[Any] = None, timeout: Optional[float] = None) -> Any:
        if timeout is None:
            timeout = self.timeout

        if not self._is_connected():
            raise ConnectionError("not connected")

        raw_id = next(self._id_counter)   # integer id — send numeric in JSON
        # build request with numeric id (important: JSON number, not string)
        req_obj = {"jsonrpc": "2.0", "id": raw_id, "method": method}
        if params is not None:
            req_obj["params"] = params

        payload = json.dumps(req_obj)
        fut = asyncio.get_event_loop().create_future()
        # store both numeric and string keys to be tolerant to server reply types
        self._pending[raw_id] = fut
        self._pending[str(raw_id)] = fut

        logger.debug("SENT RAW: %s", payload)
        try:
            await self._ws.send(payload)  # type: ignore[arg-type]
        except Exception:
            # cleanup
            self._pending.pop(raw_id, None)
            self._pending.pop(str(raw_id), None)
            raise

        try:
            result = await asyncio.wait_for(fut, timeout=timeout)
            logger.debug("Matched response for id=%s -> %r", raw_id, result)
            return result
        finally:
            # ensure cleanup of both keys
            self._pending.pop(raw_id, None)
            self._pending.pop(str(raw_id), None)

    async def send_notification(self, method: str, params: Optional[Any] = None) -> None:
        if not self._is_connected():
            raise ConnectionError("not connected")
        obj = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            obj["params"] = params
        payload = json.dumps(obj)
        logger.debug("SENT notification RAW: %s", payload)
        await self._ws.send(payload)  # type: ignore[arg-type]

    async def request_and_get_result(self, method: str, params: Optional[Any] = None, timeout: Optional[float] = None) -> Any:
        msg = await self.send_request(method, params=params, timeout=timeout)
        if "error" in msg and msg["error"] is not None:
            raise RuntimeError(f"JSON-RPC error: {msg['error']}")
        return msg.get("result")


if __name__ == "__main__":
    import sys
    async def main():
        uri = sys.argv[1] if len(sys.argv) > 1 else "ws://127.0.0.1:8000/mcp"
        async with MCPWebSocketClient(uri) as client:
            try:
                try:
                    result = await client.request_and_get_result("hello", params={"name": "world"}, timeout=5.0)
                    print("Result:", result)
                except asyncio.TimeoutError:
                    print("Request timed out")
                except Exception:
                    logger.exception("request failed")

                await client.send_notification("notify_me", params={"foo": "bar"})
                await asyncio.sleep(2.0)
            finally:
                pass

    asyncio.run(main())
